from media.vencoder import *
from media.sensor import *
from media.media import *
import time, os
import _thread
import multimedia as mm
import network  # 导入网络模块

# WiFi配置
SSID = "tplink-1"
PASSWORD = "12345678."

class RtspServer:
    def __init__(self, session_name="test", port=8554, 
                 video_type=mm.multi_media_type.media_h264, enable_audio=False):
        self.session_name = session_name
        self.video_type = video_type
        self.enable_audio = enable_audio
        self.port = port
        self.rtspserver = mm.rtsp_server()
        self.venc_chn = VENC_CHN_ID_0
        self.start_stream = False
        self.runthread_over = False

    def start(self):
        # 初始化推流相关组件
        self._init_stream()
        self.rtspserver.rtspserver_init(self.port)
        self.rtspserver.rtspserver_createsession(self.session_name, self.video_type, self.enable_audio)
        self.rtspserver.rtspserver_start()
        self._start_stream()

        # 启动推流线程
        self.start_stream = True
        _thread.start_new_thread(self._do_rtsp_stream, ())

    def stop(self):
        if not self.start_stream:
            return
        # 等待推流线程结束
        self.start_stream = False
        while not self.runthread_over:
            time.sleep(0.1)
        self.runthread_over = False

        # 停止推流并释放资源
        self._stop_stream()
        self.rtspserver.rtspserver_stop()
        self.rtspserver.rtspserver_deinit()

    def get_rtsp_url(self):
        return self.rtspserver.rtspserver_getrtspurl(self.session_name)

    def _init_stream(self):
        # 降低分辨率为640×480
        width = 640
        height = 480
        width = ALIGN_UP(width, 16)
        # 初始化摄像头
        self.sensor = Sensor()
        self.sensor.reset()
        self.sensor.set_framesize(width=width, height=height, alignment=12)
        self.sensor.set_pixformat(Sensor.YUV420SP)
        # 初始化编码器
        self.encoder = Encoder()
        self.encoder.SetOutBufs(self.venc_chn, 8, width, height)
        # 绑定摄像头与编码器
        self.link = MediaManager.link(
            self.sensor.bind_info()['src'], 
            (VIDEO_ENCODE_MOD_ID, VENC_DEV_ID, self.venc_chn)
        )
        MediaManager.init()
        # 创建编码器通道
        chnAttr = ChnAttrStr(
            self.encoder.PAYLOAD_TYPE_H264, 
            self.encoder.H264_PROFILE_MAIN, 
            width, height
        )
        self.encoder.Create(self.venc_chn, chnAttr)

    def _start_stream(self):
        self.encoder.Start(self.venc_chn)
        self.sensor.run()

    def _stop_stream(self):
        self.sensor.stop()
        del self.link
        self.encoder.Stop(self.venc_chn)
        self.encoder.Destroy(self.venc_chn)
        MediaManager.deinit()

    def _do_rtsp_stream(self):
        try:
            streamData = StreamData()
            while self.start_stream:
                os.exitpoint()
                # 获取编码后的视频流
                self.encoder.GetStream(self.venc_chn, streamData)
                # 发送视频数据到RTSP服务器
                for pack_idx in range(streamData.pack_cnt):
                    stream_data = bytes(
                        uctypes.bytearray_at(
                            streamData.data[pack_idx], 
                            streamData.data_size[pack_idx]
                        )
                    )
                    self.rtspserver.rtspserver_sendvideodata(
                        self.session_name, 
                        stream_data, 
                        streamData.data_size[pack_idx], 
                        1000
                    )
                # 释放流数据
                self.encoder.ReleaseStream(self.venc_chn, streamData)
                time.sleep(0.01)  # 轻微延时，避免资源占用过高
        except BaseException as e:
            print(f"推流异常: {e}")
        finally:
            self.runthread_over = True
            self.stop()

# 连接WiFi（参考MQTT程序的网络配置逻辑）
def connect_wifi():
    wlan = network.WLAN(network.STA_IF)  # 创建STA模式WiFi对象
    wlan.active(True)  # 激活WiFi
    if not wlan.isconnected():
        print(f"正在连接WiFi: {SSID}")
        wlan.connect(SSID, PASSWORD)  # 连接指定WiFi
        # 等待连接成功（最多等待20秒）
        wait_count = 0
        while not wlan.isconnected() and wait_count < 20:
            time.sleep(1)
            wait_count += 1
    # 检查连接状态
    if wlan.isconnected():
        ip_info = wlan.ifconfig()
        print(f"WiFi连接成功，IP地址: {ip_info[0]}")
        return wlan, ip_info[0]
    else:
        print("WiFi连接失败，请检查配置")
        return None, None

if __name__ == "__main__":
    os.exitpoint(os.EXITPOINT_ENABLE)
    
    # 先连接WiFi
    wlan, ip_addr = connect_wifi()
    if not wlan or not ip_addr:
        print("网络连接失败，无法启动推流")
    else:
        # 启动RTSP推流服务
        rtsp_server = RtspServer()
        rtsp_server.start()
        # 打印实际可访问的RTSP地址（替换0.0.0.0为设备IP）
        original_url = rtsp_server.get_rtsp_url()
        actual_url = original_url.replace("0.0.0.0", ip_addr)
        print(f"RTSP服务已启动，地址: {actual_url}")
        
        # 推流持续时间（300秒=5分钟）
        try:
            time.sleep(300)
        finally:
            # 停止推流服务
            rtsp_server.stop()
            wlan.disconnect()  # 断开WiFi
            print("推流结束，已断开网络连接")